﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using Accord.Statistics.Models.Markov;
using Accord.Statistics.Distributions.Univariate;
using Accord.Statistics.Distributions.Multivariate;
using Accord.Statistics.Models.Markov.Learning;
using Accord.Statistics.Models.Markov.Topology;
using Accord.Statistics.Distributions.Fitting;


namespace NaoController
{

    /// <summary>
    /// Used to train and use a HMM learner for skeleton data.
    /// </summary>
    class GestureLearner
    {
        public const int GESTURE_COUNT = 4;

        SequenceClassifier<MultivariateNormalDistribution> anglesClassifier = new SequenceClassifier<MultivariateNormalDistribution>(GESTURE_COUNT, new Forward(GESTURE_COUNT)/*new Ergodic(GESTURE_COUNT)*/, new MultivariateNormalDistribution(8));
        SequenceClassifier<MultivariateNormalDistribution> positionClassifier = new SequenceClassifier<MultivariateNormalDistribution>(GESTURE_COUNT, new Forward(GESTURE_COUNT)/*new Ergodic(GESTURE_COUNT)*/, new MultivariateNormalDistribution(3));

        #region Classifier Training
        /// <summary>
        /// Trains HMM on set of data
        /// </summary>
        public void train(List<KeyValuePair<HmmGesture, Dictionary<Motor, List<KeyValuePair<DateTime, float>>>>> angleData, List<KeyValuePair<HmmGesture, List<KeyValuePair<DateTime, Position>>>> positionData, bool addNoise = false)
        {
            // Increase training data size
            if (addNoise)
            {
                throw new NotImplementedException();
            }

            // Oh Lawd lambda functions
            SequenceClassifierLearning<MultivariateNormalDistribution> anglesTeacher = new SequenceClassifierLearning<MultivariateNormalDistribution>(anglesClassifier,
                modelIndex => new BaumWelchLearning<MultivariateNormalDistribution>(anglesClassifier.Models[modelIndex])
                {
                    Tolerance = 0.0001,
                    Iterations = 0,
                    FittingOptions = new NormalOptions() { Regularization = 1e-5 }
                }
            );
            SequenceClassifierLearning<MultivariateNormalDistribution> positionTeacher = new SequenceClassifierLearning<MultivariateNormalDistribution>(positionClassifier,
                 modelIndex => new BaumWelchLearning<MultivariateNormalDistribution>(positionClassifier.Models[modelIndex])
                 {
                     Tolerance = 0.0001,
                     Iterations = 0,
                     FittingOptions = new NormalOptions() { Regularization = 1e-5 }
                 }
            );

            Random rand = new Random();
            int ai = rand.Next(angleData.Count);
            Dictionary<Motor, List<KeyValuePair<DateTime, float>>> crossvalidationa = angleData[ai].Value;
            HmmGesture cvLabelA = angleData[ai].Key;
            angleData.RemoveAt(ai);
            int pi = rand.Next(positionData.Count);
            List<KeyValuePair<DateTime, Position>> crossvalidationp = positionData[pi].Value;
            HmmGesture cvLabelP = positionData[pi].Key;
            positionData.RemoveAt(pi);
            
            Console.WriteLine("Training HMMs...");
            Console.WriteLine("Trained anglesClassifier with logLikelihood: {0}", anglesTeacher.Run(convertTrainingAngles(angleData), convertAnglesGestures(angleData)));
            Console.WriteLine("Trained positionClassifier with logLikelihood: {0}", positionTeacher.Run(convertTrainingPositions(positionData), convertPositionGestures(positionData)));
            Console.WriteLine("Classifiers can classify {0} {1}", anglesClassifier.Classes, positionClassifier.Classes);
            evaluateAngles(crossvalidationa);
            Console.WriteLine("Should have evaluated to be of class {0}", cvLabelA);
            evaluatePositions(crossvalidationp);
            Console.WriteLine("Should have evaluated to be of class {0}", cvLabelP);
        }
        #endregion

        #region Classifier Evaluation
        /// <summary>
        /// Evaluates a gesture and returns what gesture the HMM thinks it is
        /// </summary>
        /// <returns>The gesture with the highest probability</returns>
        public HmmGesture evaluateAngles(Dictionary<Motor, List<KeyValuePair<DateTime, float>>> angles)
        {
            double likelihood;
            int sequenceClass = 4 + anglesClassifier.Compute(convertAngles(angles), out likelihood);
            Console.WriteLine("Tested angles sequence and evaluated as class: {0} with likelihood: {1}", (HmmGesture)(sequenceClass), likelihood);
            return (HmmGesture)sequenceClass;
        }

        /// <summary>
        /// Evaluates a gesture and returns what gesture the HMM thinks it is
        /// </summary>
        /// <returns>The gesture with the highest probability</returns>
        public HmmGesture evaluatePositions(List<KeyValuePair<DateTime, Position>> positions)
        {
            double likelihood;
            int sequenceClass = positionClassifier.Compute(convertPositions(positions), out likelihood);
            Console.WriteLine("Tested positions sequence and evaluated as class: {0} with likelihood: {1}", (HmmGesture)sequenceClass, likelihood);
            return (HmmGesture)sequenceClass;
        }
        #endregion

        #region Conversion Methods
        private int[] convertAnglesGestures(List<KeyValuePair<HmmGesture, Dictionary<Motor, List<KeyValuePair<DateTime, float>>>>> gestures)
        {
            int[] retval = new int[gestures.Count];
            for (int i = 0; i < gestures.Count; i++)
            {
                retval[i] = (int)gestures[i].Key - 4;
            }
            return retval;
        }

        private int[] convertPositionGestures(List<KeyValuePair<HmmGesture, List<KeyValuePair<DateTime, Position>>>> gestures)
        {
            int[] retval = new int[gestures.Count];
            for (int i = 0; i < gestures.Count; i++)
            {
                retval[i] = (int)gestures[i].Key;
            }
            return retval;
        }

        private Motor[] motors = {Motor.LShoulderPitch,
                                     Motor.LShoulderRoll,
                                     Motor.LElbowRoll,
                                     Motor.LElbowYaw,
                                     Motor.RShoulderPitch,
                                     Motor.RShoulderRoll,
                                     Motor.RElbowRoll,
                                     Motor.RElbowYaw};

        private double[][] convertAngles(Dictionary<Motor, List<KeyValuePair<DateTime, float>>> angles)
        {
            int tCount = angles[Motor.LElbowYaw].Count;
            int aCount = motors.Length;
            double[][] data = new double[tCount][];

            for (int i = 0; i < tCount; i++)
            {
                data[i] = new double[aCount];
                for (int j = 0; j < aCount; j++)
                {
                    data[i][j] = angles[motors[j]][i].Value;
                }
            }

            return data;
        }

        private double[][][] convertTrainingAngles(List<KeyValuePair<HmmGesture, Dictionary<Motor, List<KeyValuePair<DateTime, float>>>>> angles)
        {
            double[][][] retval = new double[angles.Count][][];
            for (int i = 0; i < angles.Count; i++)
            {
                retval[i] = convertAngles(angles[i].Value);
            }

            return retval;
        }

        private double[][] convertPositions(List<KeyValuePair<DateTime, Position>> positions)
        {
            int tCount = positions.Count;
            double[][] data = new double[tCount][];

            for (int i = 0; i < tCount; i++)
            {
                data[i] = new double[] { positions[i].Value.x, positions[i].Value.y, positions[i].Value.z };
            }

            return data;
        }

        private double[][][] convertTrainingPositions(List<KeyValuePair<HmmGesture, List<KeyValuePair<DateTime, Position>>>> positions)
        {
            double[][][] retval = new double[positions.Count][][];
            for (int i = 0; i < positions.Count; i++)
            {
                retval[i] = convertPositions(positions[i].Value);
            }

            return retval;
        }
        #endregion
    }
}
